/*
 * Copyright (c) 2015-2018, Marco Frigerio
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice, this
 *    list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright notice,
 *    this list of conditions and the following disclaimer in the documentation
 *    and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#pragma once

#include <iostream>
#include <ctime>

#include "pronto_quadruped_commons/rbd/rbd.h"
#include "pronto_quadruped_commons/rbd/utils.h"
#include "pronto_quadruped_commons/robcogen/scalar/internals.h"

#include "pronto_quadruped_commons/rbd/utils.h"

namespace pronto {
namespace robcogen {

/**
 * Test routines for the C++ code generated by RobCoGen
 */
namespace test {

namespace {

template<typename ROB>
using Scalar = typename internal::ScalarTraitsSelector< ROB >::trait::Scalar;

}
/**
 * \name Dynamics tests
 * These functions compare the output of three dynamics algorithms to verify
 * they are consistent.
 *
 * Comparison results are simply reported by printing the difference between
 * two values, which should be a zero vector.
 *
 * The tests involve the quantities of the well known dynamics equation
 * \f[ M \ddot{q} + C + G = \tau \f]
 * which are computed in various ways using the engines for inverse-dynamics,
 * forward-dynamics and the Joint Space Inertia Matrix (JSIM), using random
 * values for the input variables.
 *
 * The general test strategies are the following:
 * - (a) Inverse dynamics and JSIM:
 *    use the inverse dynamics engine to compute, separately, the gravity plus
 *    Coriolis joint forces (\f$C+G\f$) and the total joint forces
 *    (\f$\tau\f$); compute explicitly the inertia matrix (\f$M\f$) using the
 *    same joint status and then make sure that the obtained results satisfy
 *    the dynamics equation.
 *
 * - (b) Inverse and forward dynamics:
 *    use the forward dynamics engine to compute the joint accelerations,
 *    using the same position and velocity input as in test 1 and also using
 *    the inverse dynamics joint forces as input; the output joint
 *    accelerations should be equal to the input ones used with inverse
 *    dynamics. In other words, make sure that inverse and forward dynamics are
 *    actually complementary.
 *
 * - (c) Forward dynamics and JSIM inverse:
 *    compare the joint accelerations computed with the forward dynamics and
 *    with the inverse of the joint space inertia matrix \f$M^{-1}\f$; the
 *    joint velocities are set to zero, the gravity terms \f$G\f$ are
 *    computed with the inverse dynamics engine. The accelerations are
 *    computed with
 *    \f$ \ddot{q} = M^{-1} (\tau - G) \f$
 *
 * \tparam RT the class Traits for a robot, as generated by RobCoGen
 */
///@{
/**
 * Test of the dynamics of a fixed-base robot.
 *
 * This function performs the test (a) and (b) described in the general
 * comments (above).
 *
 * \param fwdDynEngine an instance of the forward dynamics class of the robot
 * \param invDynEngine an instance of the inverse dynamics class
 * \param jsim an instance of the Joint Space Inertia Matrix
 */
template<class RT>
void fixedBaseDynamics(
    typename RT::FwdDynEngine& fwdDynEngine,
    typename RT::InvDynEngine& invDynEngine,
    typename RT::JSIM&         jsim)
{
    iit::rbd::Utils::CwiseAlmostZeroOp< Scalar<RT> > approx(1E-5);
    typename RT::JointState q, qd, qdd, qdd_fd, c, g, tau, diff;
    robcogen::utils::rand_jstate<RT>(q, qd, qdd);

    invDynEngine.G_terms(g, q);
    invDynEngine.C_terms(c, q, qd);
    invDynEngine.id(tau, q, qd, RT::JointState::Zero());
    diff = tau - c - g;
    std::cout << "Inverse dynamics (consistency of the routines to compute gravity and Coriolis terms separately):" << std::endl;
    std::cout << diff.transpose().unaryExpr(approx) << std::endl << std::endl;

    invDynEngine.G_terms(g, q);
    invDynEngine.C_terms(c, q, qd);
    invDynEngine.id(tau, q, qd, qdd);
    diff = tau - (jsim(q) * qdd + c + g);
    std::cout << "Inverse dynamics and JSIM:" << std::endl;
    std::cout << diff.transpose().unaryExpr(approx) << std::endl << std::endl;

    fwdDynEngine.fd(qdd_fd, q, qd, tau);
    diff = qdd - qdd_fd;
    std::cout << "Forward dynamics:" << std::endl;
    std::cout << diff.transpose().unaryExpr(approx) << std::endl << std::endl;
}

/**
 * Test of the dynamics of a floating-base robot.
 *
 * This function performs three checks. First, it executes test (a) described
 * in the general documentation above, using the inverse dynamics routine in
 * the assumption of a fully-actuated floating-base (and the inertia matrix).
 *
 * Then it checks the consistency between the two inverse dynamics functions:
 * the one for a fully-actuated base and the regular floating-base inverse
 * dynamics (which is really an hybrid dynamics algorithm).
 *
 * Third, the function performs test (c).
 *
 * \param fwdDynEngine an instance of the forward dynamics class of the robot
 * \param invDynEngine an instance of the inverse dynamics class
 * \param jsim an instance of the Joint Space Inertia Matrix
 */
template<class RT>
void floatingBaseDynamics(
    typename RT::FwdDynEngine& fwdDynEngine,
    typename RT::InvDynEngine& invDynEngine,
    typename RT::JSIM&         jsim)
{
    iit::rbd::Utils::CwiseAlmostZeroOp< Scalar<RT> > approx(1E-10);
    typename RT::JointState q, qd, qdd, qdd_fd, tau_g, tau_c, tau, diff;
    iit::rbd::VelocityVector base_v, base_a, base_a_fd, diff_base_a, gravity;
    iit::rbd::ForceVector base_wrench_g, base_wrench_c, base_wrench, diff_wrench;

    robcogen::utils::rand_jstate<RT>(q, qd, qdd);
    iit::rbd::Utils::randomVec(base_v);
    iit::rbd::Utils::randomVec(base_a);
    iit::rbd::Utils::randomGravity(gravity);

    invDynEngine.setJointStatus(q);
    invDynEngine.G_terms_fully_actuated(base_wrench_g, tau_g, gravity);
    invDynEngine.C_terms_fully_actuated(base_wrench_c, tau_c, base_v, qd);
    invDynEngine.id_fully_actuated(base_wrench, tau, gravity, base_v, base_a, qd, qdd);
    jsim(q);

    diff = tau - (jsim.getF().transpose()*base_a + jsim.getFixedBaseBlock() * qdd + tau_c + tau_g);
    diff_wrench = base_wrench -
            (jsim.getWholeBodyInertia() * base_a +
                    jsim.getF() * qdd +
                    base_wrench_g + base_wrench_c);

    std::cout << "Fully-actuated-floating-base inverse dynamics and JSIM:" << std::endl;
    std::cout <<
            diff_wrench.transpose().unaryExpr(approx) << "   " <<
            diff.transpose().unaryExpr(approx) << std::endl << std::endl;


    // Fully-actuated inverse dynamics asking for zero base acceleration; this
    //  call returns the wrench to be applied to the base to prevent any
    //  acceleration. Such wrench is equal and opposite to the wrench acting
    //  on the base due to the motion of the actuated joints.
    invDynEngine.id_fully_actuated(base_wrench, tau, gravity, base_v, iit::rbd::VelocityVector::Zero(), qd, qdd);
    // The regular floating base inverse dynamics (ie hybrid dynamics). The
    //  resulting base acceleration (consequence of the motion of the actuated
    //  joints) should exactly correspond to the wrench to prevent acceleration
    //  that we just computed
    invDynEngine.id(tau, base_a,  gravity, base_v, qd, qdd);

    std::cout << "Fully-actuated-floating-base inverse dynamics and hybrid dynamics:" << std::endl;
    std::cout << (base_wrench + jsim.getWholeBodyInertia() * base_a)
            .transpose().unaryExpr(approx) << std::endl << std::endl;



    // Forward dynamics, check the consistency with previous results
    fwdDynEngine.fd(qdd_fd, base_a_fd, base_v, gravity, q, qd, tau);
    diff = qdd - qdd_fd;
    diff_base_a = base_a - base_a_fd;
    std::cout << "Forward dynamics:" << std::endl;
    std::cout << diff_base_a.transpose().unaryExpr(approx) << "   " <<
            diff.transpose().unaryExpr(approx) << std::endl << std::endl;
}

/**
 * Test of the inverse of the Joint Space Inertia Matrix for a fixed-base
 * robot.
 *
 * This function performs the test (c) described in the general comments
 * (above).
 *
 * \param fwdDynEngine an instance of the forward dynamics class of the robot
 * \param invDynEngine an instance of the inverse dynamics class
 * \param jsim an instance of the Joint Space Inertia Matrix
 */
template<class RT>
void fixedBaseJSIMInverse(
        typename RT::FwdDynEngine& fwdDynEngine,
        typename RT::InvDynEngine& invDynEngine,
        typename RT::JSIM&         jsim)
{
    iit::rbd::Utils::CwiseAlmostZeroOp< Scalar<RT> > approx(1E-5);
    typename RT::JointState q, qd, qdd, g, tau, diff;
    robcogen::utils::rand_jstate<RT>(q, qdd, tau);

    invDynEngine.G_terms(g, q);
    fwdDynEngine.fd(qdd, q, RT::JointState::Zero(), tau);
    jsim(q);
    jsim.computeL();
    jsim.computeInverse();
    diff = qdd - jsim.getInverse() * (tau - g);
    std::cout << "JSIM inverse:" << std::endl;
    std::cout << diff.transpose().unaryExpr(approx) << std::endl << std::endl;
}
///@}


namespace internal {

template<class ROBOT_TRAITS, bool IS_FLOATING>
struct tests_;

template<class ROBOT_TRAITS>
struct tests_<ROBOT_TRAITS, true>
{
    static void consistencyTests(
            typename ROBOT_TRAITS::FwdDynEngine& fd,
            typename ROBOT_TRAITS::InvDynEngine& id,
            typename ROBOT_TRAITS::JSIM&         jsim)
    {
        floatingBaseDynamics<ROBOT_TRAITS>(fd, id, jsim);
    }
};

template<class ROBOT_TRAITS>
struct tests_<ROBOT_TRAITS, false>
{
    static void consistencyTests(
            typename ROBOT_TRAITS::FwdDynEngine& fd,
            typename ROBOT_TRAITS::InvDynEngine& id,
            typename ROBOT_TRAITS::JSIM&         jsim)
    {
        fixedBaseDynamics<ROBOT_TRAITS>(fd, id, jsim);
        fixedBaseJSIMInverse<ROBOT_TRAITS>(fd, id, jsim);
    }
};

}



template<class ROBOT_TRAITS>
void consistencyTests(
        typename ROBOT_TRAITS::InertiaProperties& ip,
        typename ROBOT_TRAITS::MotionTransforms& xm,
        typename ROBOT_TRAITS::ForceTransforms&  xf)
{
    typename ROBOT_TRAITS::InvDynEngine id  (ip, xm);
    typename ROBOT_TRAITS::FwdDynEngine fd  (ip, xm);
    typename ROBOT_TRAITS::JSIM         jsim(ip, xf);

    internal::tests_<ROBOT_TRAITS, ROBOT_TRAITS::floating_base>::consistencyTests(fd,id,jsim);
}

template<class ROBOT_TRAITS>
void consistencyTests(
        typename ROBOT_TRAITS::FwdDynEngine& fd,
        typename ROBOT_TRAITS::InvDynEngine& id,
        typename ROBOT_TRAITS::JSIM&         jsim)
{
    internal::tests_<ROBOT_TRAITS, ROBOT_TRAITS::floating_base>::consistencyTests(fd,id,jsim);
}

}
}
}
